import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt


# 生成圆形图像的函数（使用PyTorch）
def generate_circle_image(size=64):
    image = torch.zeros((1, size, size))  # 使用 PyTorch 创建空白图像
    center = size // 2
    radius = size // 4
    for y in range(size):
        for x in range(size):
            if (x - center) ** 2 + (y - center) ** 2 <= radius ** 2:
                image[0, y, x] = 1  # 在圆内的点设置为白色
    return image


# 生成方形图像的函数（使用PyTorch）
def generate_square_image(size=64):
    image = torch.zeros((1, size, size))  # 使用 PyTorch 创建空白图像
    padding = size // 4
    image[0, padding:size - padding, padding:size - padding] = 1  # 设置方形区域为白色
    return image


# 自定义数据集：圆形和方形图像
class ShapeDataset(Dataset):
    def __init__(self, num_samples=1000, size=64):
        self.num_samples = num_samples
        self.size = size
        self.data = []
        # 生成数据：一半是圆形图像，一半是方形图像
        for i in range(num_samples // 2):
            self.data.append(generate_circle_image(size))
            self.data.append(generate_square_image(size))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx].float()  # 直接返回 PyTorch Tensor 格式的数据


# VAE模型定义
class VAE(nn.Module):
    def __init__(self, latent_dim=2):
        super(VAE, self).__init__()
        self.latent_dim = latent_dim

        # 编码器
        self.fc1 = nn.Linear(64 * 64, 400)
        self.fc21 = nn.Linear(400, latent_dim)  # 均值
        self.fc22 = nn.Linear(400, latent_dim)  # 方差

        # 解码器
        self.fc3 = nn.Linear(latent_dim, 400)
        self.fc4 = nn.Linear(400, 64 * 64)

    def encode(self, x):
        h1 = torch.relu(self.fc1(x.view(-1, 64 * 64)))
        return self.fc21(h1), self.fc22(h1)  # 返回均值和方差

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        h3 = torch.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h3)).view(-1, 1, 64, 64)  # 重构图像

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar


# 损失函数：重构误差 + KL 散度
def loss_function(recon_x, x, mu, logvar):
    BCE = nn.functional.binary_cross_entropy(recon_x.view(-1, 64 * 64), x.view(-1, 64 * 64), reduction='sum')
    # KL 散度
    return BCE + 0.5 * torch.sum(torch.exp(logvar) + mu ** 2 - 1 - logvar)


# 设置超参数
batch_size = 128
epochs = 10
latent_dim = 2
learning_rate = 1e-3

# 数据加载
train_loader = DataLoader(ShapeDataset(num_samples=2000), batch_size=batch_size, shuffle=True)

# 创建模型和优化器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = VAE(latent_dim).to(device)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)


# 训练模型
def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, data in enumerate(train_loader):
        data = data.to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()

        if batch_idx % 100 == 0:
            print(
                f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item() / len(data):.6f}')

    print(f'Train Epoch: {epoch} Average loss: {train_loss / len(train_loader.dataset):.4f}')

# 测试并显示一些真实图像和生成的图像
def test():
    model.eval()
    with torch.no_grad():
        # 获取一批真实的图像（原始图像）
        real_images = next(iter(train_loader))[:64]  # 只取前64个图像
        real_images = real_images.cpu().numpy()

        # 从潜在空间随机生成一些样本
        sample = torch.randn(64, latent_dim).to(device)
        generated_images = model.decode(sample).cpu().numpy()

        # 显示真实图像和生成的图像，分别标明
        fig, axes = plt.subplots(8, 8, figsize=(8, 8))
        axes = axes.flatten()

        for i in range(64):
            if i < 32:  # 前32个显示真实图像
                axes[i].imshow(real_images[i].squeeze(), cmap='gray')
                axes[i].set_title('Real', fontsize=8)
            else:  # 后32个显示生成图像
                axes[i].imshow(generated_images[i - 32].squeeze(), cmap='gray')
                axes[i].set_title('Generated', fontsize=8)

            axes[i].axis('off')

        plt.tight_layout()
        plt.show()

# 训练模型
for epoch in range(1, epochs + 1):
    train(epoch)

# 训练完成后，显示生成的图像
test()
